from __future__ import division
from math import *
from numpy import *
from Lists import Sat_M, Sat_Rad, Sat_r, Sat_v, r, m, sigma, deltaA, w, rhill, Sat_e, Torque_to_disk, Sat_tstep, Laplace_Coefficient, DLaplace_Coefficient, Sat_w, Sat_r_default, R
from Constants import rho_sat, G, year, R_Mars
from Init_Cond import N, r_I, deltar, m_pdisk, r_pdisk, deltar, r_F
import Timestep
import dM_dt
import sys
import Plot
import Save
import Timestep

Gamma_S_in = []
Gamma_S_out = []
T_r = []
Sat_time = []
Gamma_S = []

def construct(i_rigid, M):
#initial list of satellites
    x = r[i_rigid]      #location at edge of disk
    while x < r[int(N-1)]:    #run until the distance is larger than the extent of the bins for the disk
        Sat_M.append(0.0)
        Sat_r.append(x)         #put a placeholder for a satellite at this location
        Sat_r_default.append(x)
        Sat_v.append(sqrt(G*M/x))
        Sat_e.append(0.01)
        Sat_time.append(0.0)
        Sat_w.append(sqrt(G*M/x**3.))
        # Sat_v.append(0.0)
        Sat_Rad.append(r_pdisk) #inital satellite size is that of a disk particle
        rhill.append(x*(m_pdisk/(M*3.))**(1./3.))   #hill radius for disk particle at this location
        hill = 4.0*x*(m_pdisk/(M*3.))**(1./3.)  #4*hill radius for a disk particle at location x
        x = x + hill    #advance location to 10*hill (like embryo spacing)






def grow(i_roche, M, deltat, t):

#Accreting New satellites
    x = sum(m[i_roche:int(N)])
    if x > 0.0:         #is there mass in the disk outside the roche limit?
        for i in range(len(Sat_M)):
            if Sat_M[i] == 0.0:
                loc = int(round((Sat_r[i] - r_I)/deltar - 0.5))     #what bin of the disk does the satellite's location correspond to?
            # if loc > len(r)-1:
                # print('grow',Sat_r[i]/R_Mars)
                # continue
                if m[loc] > 1*m_pdisk and Sat_M[i] == 0.0 and Sat_r[i] > r[i_roche]:    #if that bin has enough mass, and the satellite is not turned on, build a satellite
                    Sat_M[i] = 1*m_pdisk
                    Sat_Rad[i] = (3.*Sat_M[i]/(4.*pi*rho_sat))**(1./3.)
                    Sat_v[i] = sqrt(G*M/Sat_r[i])
                    rhill[i] = Sat_r[i]*(Sat_M[i]/(M*3.))**(1./3.)
                    Sat_time[i] = 0.0
                    m[loc] = m[loc] - Sat_M[i]
                    sigma[loc] = m[loc]/deltaA[loc]

    #Accreting mass onto existing satellites
    for i in range(len(Sat_M)):
        if Sat_M[i] >= m_pdisk:

            inner = int(round((Sat_r[i] - rhill[i] - r_I)/deltar - 0.5))     #what bin of the disk does the satellite's inner feeding zone correspond to?
            outer = int(round((Sat_r[i] + rhill[i] - r_I)/deltar - 0.5))     #what bin of the disk does the satellite's outer feeding zone correspond to?

            if inner < 0:   #can't eat inside the disk bound
                inner = 0

            if outer > int(N-1):    #can't eat outside the disk bound!
                outer = int(N-1)

            m_avail = sum(m[inner:outer+1]) #I'm hungry! What's there to eat?!
            if m_avail > 0.0:


                A = 4.5
                B = 1.5
                conversion = 1.68741e-18
                P = A*(Sat_M[i]*conversion)**B
                Q = A*(Sat_M[i]*conversion + deltat/year/2.*P)**B
                R = A*(Sat_M[i]*conversion + deltat/year/2.*Q)**B
                S = A*(Sat_M[i]*conversion + deltat/year*R)**B
                msat = Sat_M[i] + deltat/year/6.*(P + 2.*Q + 2.*R + S)/conversion


                diff = msat - Sat_M[i]
                if 0.0 <= diff <= m_avail:
                    Sat_M[i] = Sat_M[i] + diff
                    Sat_Rad[i] = (3.*Sat_M[i]/(4.*pi*rho_sat))**(1./3.)
                    rhill[i] = Sat_r[i]*(Sat_M[i]/(M*3.))**(1./3.)


                    for j in range(inner,outer+1):
                        m[j] = m[j] - diff*m[j]/m_avail
                        sigma[j] = m[j]/deltaA[j]


                elif diff > m_avail:
                    Sat_M[i] = Sat_M[i] + m_avail
                    Sat_Rad[i] = (3.*Sat_M[i]/(4.*pi*rho_sat))**(1./3.)
                    rhill[i] = Sat_r[i]*(Sat_M[i]/(M*3.))**(1./3.)

                    for j in range(inner, outer+1):
                        m[j] = 0.0
                        sigma[j] = 0.0



def kill():
    for i in range(len(Sat_M)):
        if Sat_M[i] != 0.0:
            x = i

    if x < int(N-1):
        for i in range(x+1):
            if i != x and Sat_M[i] != 0.0:
                Sat_M[i] = 0.0
            elif i == x:
                print(Sat_r[i]/R_Mars)
    else:
        for i in range(len(Sat_M)):
            if i == len(Sat_M):
                print(Sat_r[i]/R_Mars)
            elif i != len(Sat_M) and Sat_M[i] != 0.0:
                Sat_M[i] = 0.0




#merging
def merge(M, last):

    for i in range(last):
        if Sat_M[i] == 0.0 or i == last-1:##If there is no satellite, or we're at the last one, don't bother doing anything and move on
        # if i == len(Sat_M):##If there is no satellite, or we're at the last one, don't bother doing anything and move on

            continue

        for j in range((i+1), last):  #Starting at the satellite adjacent to i, run through the rest of the list
            if Sat_M[j] == 0.0:
                continue        #no satellite at j, move on

            elif 2.0*rhill[i] >= abs(Sat_r[i] - Sat_r[j]) or 2.0*rhill[j] >= abs(Sat_r[i] - Sat_r[j]):    #are the distances between the satellites less than either of the satellites' feeding zone?  4 was chosen arbitrarily

                sat_loc = (Sat_M[i]*Sat_r[i] + Sat_M[j]*Sat_r[j])/(Sat_M[i] + Sat_M[j])    #new satellite location

                if sat_loc >= Sat_r_default[i+1]:       #is new location between i and j?
                    a = i
                    while a < j and sat_loc >= Sat_r_default[a+1]:  #find location for new satellite
                        a = a + 1       # a !+ i
                    Sat_r[a] = sat_loc

                    if a == j:
                        Sat_M[a] = Sat_M[i] + Sat_M[j]
                        Sat_Rad[a] = ((3.*Sat_M[a]/(4.*pi*rho_sat))**(1./3.))
                        rhill[a] = Sat_r[a]*(Sat_M[a]/(M*3.))**(1./3.)
                        Sat_v[a] = sqrt(G*M/Sat_r[a])
                        # Sat_time[a] = (Sat_time[i] + Sat_time[j])/2.0
                        Sat_time[a] = max(Sat_time[i], Sat_time[j])
                        Sat_M[i] = 0.0  #turn off satellite i
                        Sat_time[i] = 0.0
                        Sat_r[i] = Sat_r_default[i]
                        Sat_v[i] = sqrt(G*M/Sat_r[i])
                        Sat_e[i] = 0.01

                    else:
                        Sat_M[a] = Sat_M[i] + Sat_M[a] + Sat_M[j]  #new satellite mass.  (includes location a in case there's a small satellite at a (could i and a could be within hill of j)
                        Sat_Rad[a] = ((3.*Sat_M[a]/(4.*pi*rho_sat))**(1./3.))
                        rhill[a] = Sat_r[a]*(Sat_M[a]/(M*3.))**(1./3.)
                        Sat_time[a] = max(Sat_time[i:j])
                        Sat_v[a] = sqrt(G*M/Sat_r[a])
                        Sat_e[a] = mean(Sat_e[i:j])
                        Sat_M[i] = 0.0  #turn off satellite i
                        Sat_time[i] = 0.0
                        Sat_r[i] = Sat_r_default[i]
                        Sat_v[i] = sqrt(G*M/Sat_r[i])
                        Sat_e[i] = 0.01
                        Sat_M[j] = 0.0  #turn off satellite j
                        Sat_time[j] = 0.0
                        Sat_r[j] = Sat_r_default[j]
                        Sat_v[j] = sqrt(G*M/Sat_r[j])
                        Sat_e[j] = 0.01


                    break   #I've turned off satellite i, so this current for loop does not make sense; move on



                else:   #new location is at bin i
                    Sat_r[i] = sat_loc
                    Sat_M[i] = Sat_M[i] + Sat_M[j]  #new satellite mass
                    Sat_Rad[i] = ((3.*Sat_M[i]/(4.*pi*rho_sat))**(1./3.))
                    rhill[i] = Sat_r[i]*(Sat_M[i]/(M*3.))**(1./3.)
                    Sat_v[i] = sqrt(G*M/Sat_r[i])
                    Sat_M[j] = 0.0
                    Sat_time[j] = 0.0
                    Sat_r[j] = Sat_r_default[j]
                    Sat_v[j] = sqrt(G*M/Sat_r[j])
                    Sat_e[j] = 0.01





def A(ratio, a, m, switch): #switch marks whether the satellite is interior or exterior to the disk bin
    alpha = int(round(ratio/0.012,0))
    # print(alpha)
    if alpha >= 84:
        alpha = 83
    if switch == 0:
        return -G*Sat_M[a]/(2.*Sat_r[a])*(2.*m*Laplace_Coefficient[m][alpha] + alpha*DLaplace_Coefficient[m][alpha])
    else:
        return -G*Sat_M[a]/(2.*Sat_r[a])*(2.*m*Laplace_Coefficient[m][alpha] - alpha*DLaplace_Coefficient[m][alpha])


# evolving
def evolve(t, M, deltat, R_Embryo, r_sync, i_rigid, i_roche, last):
    del Sat_tstep[:]
    a = 0
    f = 10
    switch = 0
    k_2 = 0.148
    Q = 88.
    k = 100
    l = 100
    while a < len(Sat_M):   #run through all satellites

        if Sat_M[a] > 0.0:  #if a satellite exists
            del T_r[:]
            del Gamma_S[:]

            F_me = 0.0

            # first mode is done alone for outer part of disk (inner part of disk would be at the location of Mars
            mode = 1
            y = (1. + 1./mode)**(2./3.)*Sat_r[a]   #resonance locations
            j = int(round((y - r_I)/deltar - .5))   #disk location of resonance

            if j < int(N-1):  #if within bin boundaries..
                # Torque = f*mode*(mode - 1)*sigma[j]*r[j]**4.*w[j]**2.*(Sat_M[a]/M)**2.    #torque exerted on satellite by bin j
                Torque = (mode*G*Sat_M[a])**2.*sigma[j]/(w[j]**2.*r[j]**2.)
                del_m = deltat*Torque/(R[j]*w[j])       #change in mass to LR
                if del_m > m[j-1]:    #am I trying to move more mass than is available in the bin it comes from?
                    del_m = m[j-1]
                    Torque = R[j]*w[j]*del_m/deltat  #scale the torque to equal that amount moved from LR bin

                m[j] = m[j] + del_m
                sigma[j] = m[j]/deltaA[j]
                m[j+1] = m[j+1] - del_m
                sigma[j+1] = m[j+1]/deltaA[j+1]

                Gamma_S.append(-1.*Torque)  #Torque exerted on the satellite is equal and opposite to the torque the satellite exerts on the disk

            
            mode = 2
            
            while abs((1 - 1/mode)**(2./3.)*Sat_r[a] - (1 - 1/(mode + 1))**(2./3.)*Sat_r[a]) > deltar:  #while the distance between Lindblad Resonances is greater than bin width

                x = (1. - 1./mode)**(2./3.)*Sat_r[a]   #resonance locations
                i = int(round((x - r_I)/deltar - .5))   #bin ID of resonance location

                if 0 < i < int(N):

                    Torque = -1.*(mode*G*Sat_M[a])**2.*sigma[i]/(w[i]**2.*r[i]**2.)
                    del_m = deltat*Torque/(R[i]*w[i])       #change in mass to LR
                    if abs(del_m) > m[i]:    #am I trying to move more mass than is available?
                        del_m = -1.*m[i]    #mass change is negative, so keep it that way
                        Torque = R[i]*w[i]*del_m/deltat  #scale the torque to equal that amount moved from LR bin
                    m[i] = m[i] + del_m
                    sigma[i] = m[i]/deltaA[i]
                    m[i-1] = m[i-1] - del_m
                    sigma[i-1] = m[i-1]/deltaA[i-1]

                    Gamma_S.append(-1.*Torque)  #Torque exerted on the satellite is equal and opposite to the torque the satellite exerts on the disk

                y = (1. + 1./mode)**(2./3.)*Sat_r[a]   #resonance locations
                j = int(round((y - r_I)/deltar - .5))
            
                if 0 < j < int(N-1):
                    Torque = (mode*G*Sat_M[a])**2.*sigma[j]/(w[j]**2.*r[j]**2.)
                    del_m = deltat*Torque/(R[j]*w[j])       #change in mass to LR
                    if del_m > m[j-1]:    #am I trying to move more mass than is available in the bin it comes from?
                        del_m = m[j-1]
                        Torque = R[j]*w[j]*del_m/deltat  #scale the torque to equal that amount moved from LR bin

                    m[j] = m[j] + del_m
                    sigma[j] = m[j]/deltaA[j]
                    m[j+1] = m[j+1] - del_m
                    sigma[j+1] = m[j+1]/deltaA[j+1]

                    Gamma_S.append(-1.*Torque)  #Torque exerted on the satellite is equal and opposite to the torque the satellite exerts on the disk
            
                mode = mode + 1

            # ****
            Torque = 1.*sum(Gamma_S)

            n = (G*M/Sat_r[a]**3.)**0.5
            if Sat_r[a] < r_sync:
                da_dt = 2.*sqrt(Sat_r[a])*Torque/(Sat_M[a]*sqrt(G*M)) - 3.*k_2*n*Sat_M[a]*R_Embryo**5./(Q*Sat_r[a]**4.*M)   #change to orbital radius, planetary tides pull in
            else:
                da_dt = 2.*sqrt(Sat_r[a])*Torque/(Sat_M[a]*sqrt(G*M)) + 3.*k_2*n*Sat_M[a]*R_Embryo**5./(Q*Sat_r[a]**4.*M)   #change to orbital radius, planetary tides push out
            de_dt = 57.*k_2*n*Sat_M[a]*R_Embryo**5./(Q*M*Sat_r[a]**5.)*Sat_e[a] - 21*0.01*n*M*Sat_Rad[a]**5./(2.*Sat_M[a]*Sat_r[a]**5.)*Sat_e[a] + F_me


            Sat_r[a] = Sat_r[a] + deltat*da_dt      #new satellite distance
            Sat_e[a] = Sat_e[a] + deltat*de_dt
            rhill[a] = Sat_r[a]*(Sat_M[a]/(M*3.))**(1./3.)


            if a == (len(Sat_M)-1):     #if you're in the last satellite bin, there's no one further out and your distance can be as great as you want
                break
            while a < len(Sat_M) and Sat_r[a] > Sat_r_default[a+1]: #is satellite location outside its default bin (away from central body)? (a < len keeps code from running out of list)
                if Sat_M[a+1] == 0.0:   #if there is not a satellite in next bin, simply turn off a and turn on a+1
                    Sat_r[a+1] = Sat_r[a]
                    Sat_M[a+1] = Sat_M[a]
                    rhill[a+1] = rhill[a]
                    Sat_Rad[a+1] = Sat_Rad[a]
                    Sat_time[a+1] = Sat_time[a]
                    Sat_M[a] = 0.0
                    Sat_time[a] = 0.0
                    Sat_r[a] = Sat_r_default[a]
                    Sat_v[a] = sqrt(G*M/Sat_r[a])
                    Sat_e[a] = 0.01
                    a = a + 1
                else:   #next satellite has mass, merge them
                    sat_loc = (Sat_M[a]*Sat_r[a] + Sat_M[a+1]*Sat_r[a+1])/(Sat_M[a] + Sat_M[a+1])    #new satellite location
                    # if sat_loc >= Sat_r_default[a+1]:   #is new location in bin a+1?
                    Sat_r[a+1] = sat_loc
                    Sat_M[a+1] = Sat_M[a] + Sat_M[a+1]  #new satellite mass
                    # Sat_time[a+1] = (Sat_time[a] + Sat_time[a+1])/2.
                    Sat_time[a+1] = max(Sat_time[a], Sat_time[a+1])
                    Sat_Rad[a+1] = ((3.*Sat_M[a+1]/(4.*pi*rho_sat))**(1./3.))
                    rhill[a+1] = Sat_r[a+1]*(Sat_M[a+1]/(M*3.))**(1./3.)
                    Sat_v[a+1] = sqrt(G*M/Sat_r[a+1])
                    Sat_e[a+1] = (Sat_e[a] + Sat_e[a+1])/2.
                    Sat_M[a] = 0.0  #merged, turn off satellite a
                    Sat_r[a] = Sat_r_default[a]
                    Sat_v[a] = sqrt(G*M/Sat_r[a])
                    a = a + 1

                if a == len(Sat_r)-1:
                    break

            while a > 0 and Sat_r[a] < Sat_r_default[a]: #is satellite location outside its default bin (towards central body)? (a>0 keeps code from blowing up)
                if Sat_M[a-1] == 0.0:   #if there is not a satellite in next bin, simply turn off a and turn on a-1
                    Sat_r[a-1] = Sat_r[a]
                    Sat_M[a-1] = Sat_M[a]
                    rhill[a-1] = rhill[a]
                    Sat_Rad[a-1] = Sat_Rad[a]
                    Sat_time[a-1] = Sat_time[a]
                    Sat_M[a] = 0.0
                    Sat_time[a] = 0.0
                    Sat_r[a] = Sat_r_default[a]
                    Sat_v[a] = sqrt(G*M/Sat_r[a])
                    Sat_e[a] = 0.01
                    a = a - 1
                else:   #next satellite has mass, merge them
                    sat_loc = (Sat_M[a]*Sat_r[a] + Sat_M[a-1]*Sat_r[a-1])/(Sat_M[a] + Sat_M[a-1])    #new satellite location
                    # if sat_loc >= Sat_r_default[a+1]:   #is new location in bin a+1?
                    Sat_r[a-1] = sat_loc
                    Sat_M[a-1] = Sat_M[a] + Sat_M[a-1]  #new satellite mass
                    # Sat_time[a-1] = (Sat_time[a] + Sat_time[a-1])/2.
                    Sat_time[a-1] = max(Sat_time[a], Sat_time[a-1])
                    Sat_Rad[a-1] = ((3.*Sat_M[a-1]/(4.*pi*rho_sat))**(1./3.))
                    rhill[a-1] = Sat_r[a-1]*(Sat_M[a-1]/(M*3.))**(1./3.)
                    Sat_v[a-1] = sqrt(G*M/Sat_r[a-1])
                    Sat_e[a-1] = (Sat_e[a] + Sat_e[a-1])/2.
                    Sat_M[a] = 0.0  #merged, turn off satellite a
                    Sat_time[a] = 0.0
                    Sat_r[a] = Sat_r_default[a]
                    Sat_v[a] = sqrt(G*M/Sat_r[a])
                    a = a - 1

            if Sat_r[a] <= r[i_rigid]:
                Rigid_Roche = r[i_rigid]
                Roche = r[i_roche]

                Save.f(t, M, Rigid_Roche, Roche, r_sync)
                Plot.sat(t, M, R_Embryo, r_sync, Roche, last, Rigid_Roche)

                m[i_rigid] = m[i_rigid] + 0.4*Sat_M[a]
                m[i_rigid-1] = m[i_rigid-1] + 0.2*Sat_M[a]
                m[i_rigid+1] = m[i_rigid+1] + 0.2*Sat_M[a]
                m[i_rigid-2] = m[i_rigid-2] + 0.1*Sat_M[a]
                m[i_rigid+2] = m[i_rigid+2] + 0.1*Sat_M[a]
                sigma[i_rigid] = m[i_rigid]/deltaA[i_rigid]
                sigma[i_rigid-1] = m[i_rigid-1]/deltaA[i_rigid-1]
                sigma[i_rigid+1] = m[i_rigid+1]/deltaA[i_rigid+1]
                sigma[i_rigid-2] = m[i_rigid-2]/deltaA[i_rigid-2]
                sigma[i_rigid+2] = m[i_rigid+2]/deltaA[i_rigid+2]

                Sat_M[a] = 0.0
                Sat_r[a] = Sat_r_default[a]


        a = a + 1 #advance to next satellite (while loop)

def final():
    x = len(Sat_M)
    last = x-1
    for a in range(x):
        if Sat_M[a] > 0.0:
            last = a

    if last == x:
        return last
    else:
        return last + 1


